// Copyright (C) 2019. Huawei Technologies Co., Ltd. All rights reserved.

// Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
// WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

#ifndef _H_RESHAPEREDUCEMEANOPTIMIZER
#define _H_RESHAPEREDUCEMEANOPTIMIZER

#include "OPOptimizer.hpp"

class ReshapeReduceMeanOptimizer : public OPOptimizer {
    bool optimize(ModelSpec *spec) override
    {
        bool hasOptimized = false;
        for (int i = 0; i < spec->num_operator_specs; i++) {
            if (OT_Reshape == spec->ops[i].type) {
                int prevOpIndex = i;
                std::vector<std::pair<int, int>> nextOpIndexes = searchOperatorIndexByInput(spec, spec->ops[prevOpIndex].output_tensors_name[0], prevOpIndex + 1, spec->num_operator_specs);
		        ReshapeParamSpec reshapePs = spec->ops[prevOpIndex].ps.reshape_spec;
		        if (reshapePs.num_shape != 3 || reshapePs.shape[1] != 1) {
		            continue;
		        }
                if (nextOpIndexes.size() != 1) {
                    continue;
		        }
		        if (spec->ops[nextOpIndexes[0].first].type != OT_Reduction) {
	                continue;
		        }
		        ReductionParamSpec reductionPs = spec->ops[nextOpIndexes[0].first].ps.reduction_spec;
		        if (reductionPs.axes[0] == 1 && reductionPs.keep_dim == false && reductionPs.mode == REDUCTION_MEAN) {		
                    setOperatorInvalid(spec, prevOpIndex, true);
		            setOperatorInvalid(spec, nextOpIndexes[0].first, true);
                    hasOptimized = true;
		        }
            }
        }
        return hasOptimized;
    }
};
#endif
